# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: chensili

from recordclass import recordclass

from coin.base.book.types import TradeSide

CrossStats = recordclass('CrossStats',
                         [
                             'product',
                             'normal_count',
                             'avg_normal_size',
                             'black_count',
                             'avg_black_size',
                             'buy_sweep_count',
                             'avg_buy_sweep_size',
                             'sell_sweep_count',
                             'avg_sell_sweep_size',
                             'between_book_trade',
                             'outside_bid_ask_trade'
                         ])

BetweenBookTrade = recordclass(
    'BetweenBookTrade',
    [
        'begin',  # Begin time of max case.
        'time_delta',  # Time duration of max case.
        'max_between_book_trade_count',  # Max of trades between two book.
        'total_between_book_trade_count',
        'stats_count'
    ])

OutsideBidAskTrade = recordclass(
    'OutsideBidAskTrade',
    [
        'begin',
        'time_delta',
        'max_outside_bid_ask_trade_count',  # Max of trades outside bid ask between two book.
        'total_outside_bid_ask_trade_count',
        'stats_count'
    ])

StatsCounter = recordclass('StatsCounter', ['begin', 'end', 'count'])


def gen_empty_CrossStats():
  return CrossStats(product=None,
                    normal_count=0,
                    avg_normal_size=0,
                    black_count=0,
                    avg_black_size=0,
                    buy_sweep_count=0,
                    avg_buy_sweep_size=0,
                    sell_sweep_count=0,
                    avg_sell_sweep_size=0,
                    between_book_trade=BetweenBookTrade(0, 0, 0, 0, 0),
                    outside_bid_ask_trade=OutsideBidAskTrade(0, 0, 0, 0, 0))


def count_between_book_trade(trade, counter):
  if counter.begin == 0:
    counter.begin = trade.timestamp
  counter.end = trade.timestamp
  counter.count += 1


def calculate_between_book_trade(cross_stats, counter):
  # Calculate only when there are trades between 2 books.
  if counter.begin > 0:
    if (counter.count > cross_stats.between_book_trade.max_between_book_trade_count):
      cross_stats.between_book_trade.begin = counter.begin
      cross_stats.between_book_trade.time_delta = counter.end - counter.begin
      cross_stats.between_book_trade.max_between_book_trade_count = counter.count
    cross_stats.between_book_trade.total_between_book_trade_count += counter.count
    cross_stats.between_book_trade.stats_count += 1


def count_outside_bid_ask_trade(trade, counter, bid_price, ask_price):
  if counter.begin == 0:
    counter.begin = trade.timestamp
  counter.end = trade.timestamp
  # If price is the same as bid or ask, still counts.
  if trade.price >= ask_price or trade.price <= bid_price:
    counter.count += 1


def calculate_outside_bid_ask_trade(cross_stats, counter):
  # Calculate only when there are beyong bid ask range trades between 2 books.
  if counter.begin > 0:
    if (counter.count > cross_stats.outside_bid_ask_trade.max_outside_bid_ask_trade_count):
      cross_stats.outside_bid_ask_trade.begin = counter.begin
      cross_stats.outside_bid_ask_trade.time_delta = counter.end - counter.begin
      cross_stats.outside_bid_ask_trade.max_outside_bid_ask_trade_count = counter.count
    if counter.count > 0:
      cross_stats.outside_bid_ask_trade.total_outside_bid_ask_trade_count += counter.count
      cross_stats.outside_bid_ask_trade.stats_count += 1


class CrossStatsCalculator(object):
  def __init__(self, cross_stats):
    self._cross_stats = cross_stats
    self._prev_book = None
    self._prev_trade = None
    self._between_book_trade_counter = StatsCounter(0, 0, 0)
    self._outside_bid_ask_trade_counter = StatsCounter(0, 0, 0)

  def set_product(self, product):
    self.cross_stats.product = product

  @property
  def cross_stats(self):
    return self._cross_stats

  def update_by_book(self, book):
    self._prev_book = book
    calculate_between_book_trade(self.cross_stats, self._between_book_trade_counter)
    calculate_outside_bid_ask_trade(self.cross_stats, self._outside_bid_ask_trade_counter)
    # Reset counter.
    self._between_book_trade_counter = StatsCounter(0, 0, 0)
    self._outside_bid_ask_trade_counter = StatsCounter(0, 0, 0)

  def update_by_trade(self, trade):
    cross_stats = self.cross_stats
    if self._prev_book is not None:
      prev_book = self._prev_book
      try:
        bid0_price = prev_book.bid0().price
        ask0_price = prev_book.ask0().price
      except (AttributeError, IndexError):
        return

      if prev_book.bid0().price < trade.price < prev_book.ask0().price:
        avg_size = cross_stats.avg_black_size
        cnt = cross_stats.black_count
        avg_size = (avg_size * cnt + trade.qty) / (cnt + 1)
        cross_stats.avg_black_size = avg_size
        cross_stats.black_count += 1
      elif trade.price > prev_book.ask0().price:
        avg_size = cross_stats.avg_buy_sweep_size
        cnt = cross_stats.buy_sweep_count
        avg_size = (avg_size * cnt + trade.qty) / (cnt + 1)
        cross_stats.avg_buy_sweep_size = avg_size
        cross_stats.buy_sweep_count += 1
      elif trade.price < prev_book.bid0().price:
        avg_size = cross_stats.avg_sell_sweep_size
        cnt = cross_stats.sell_sweep_count
        avg_size = (avg_size * cnt + trade.qty) / (cnt + 1)
        cross_stats.avg_sell_sweep_size = avg_size
        cross_stats.sell_sweep_count += 1
      else:
        avg_size = cross_stats.avg_normal_size
        cnt = cross_stats.normal_count
        avg_size = (avg_size * cnt + trade.qty) / (cnt + 1)
        cross_stats.avg_normal_size = avg_size
        cross_stats.normal_count += 1
      count_between_book_trade(trade, self._between_book_trade_counter)
      count_outside_bid_ask_trade(trade,
                                  self._outside_bid_ask_trade_counter,
                                  prev_book.bid0().price,
                                  prev_book.ask0().price)
